export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export DATA_DIR="/data_nvme0n1p2/hemuhui/datasets/cat"

DEVICE_ID=$1
output_dir_name=$2

output_dir="../checkpoint/textual_inversion/"
output_dir="${output_dir}${output_dir_name}"

mkdir ${output_dir}
if [ $? -ne 0 ]; then
  exit 1
fi

python examples/textual_inversion/textual_inversion.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --train_data_dir=$DATA_DIR \
  --learnable_property="object" \
  --placeholder_token="<cat-toy>" --initializer_token="toy" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=3000 \
  --learning_rate=5.0e-04 --scale_lr \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --output_dir=${output_dir} | tee ${output_dir}/train.log